# 你想让你的对象支持上下文管理协议(with 语句)。

# 为了让一个对象兼容with 语句，你需要实现 __ente__() 和 __exit__() 方法。

from socket import socket, AF_INET, SOCK_STREAM

class LazyConnection:
    def __init__(self, address, family = AF_INET, type = SOCK_STREAM):
        self.address = address
        self.family = family
        self.type = type
        self.sock = None

    def __enter__(self):
        if self.sock is not None:
            raise RuntimeError('Already connected')
        self.sock = socket(self.family, self.type)
        self.sock.connect(self.address)
        return self.sock

    # 不管 with 代码块中发生什么，__exit__()都会执行，就算代码块中发生了异常也是一样的。
    # 事实上，__exit__() 方法的第三个参数包含了异常类型、异常值和追溯信息(如果有的话)。
    # __exit__() 方法能自己决定怎样利用这个异常信息，或者忽略它并返回一个None 值。
    # 如果 __exit__() 返回True ，那么异常会被清空，就好像什么都没发生一样，with 语句后面的程序继续在正常执行。
    def __exit__(self, exc_ty, exc_val, tb):
        self.sock.close()
        self.sock = None

from functools import partial
conn = LazyConnection(('www.baidu.com', 80))
with conn as s:
    s.send(b'GET / HTTP/1.0\r\n')
    s.send(b'Host: www.baidu.com\r\n')
    s.send(b'\r\n')
    resp = b''.join(iter(partial(s.recv, 8192), b''))
    print(resp.decode('utf-8'))
